import sys
import os
import time
import tensorflow as tf
import numpy as np

hifi0_model = tf.keras.models.load_model('tfmodel/hifigan0')
hifi1_model = tf.keras.models.load_model('tfmodel/hifigan1')
hifi2_model = tf.keras.models.load_model('tfmodel/hifigan2')
hifi3_model = tf.keras.models.load_model('tfmodel/hifigan3')
hifi4_model = tf.keras.models.load_model('tfmodel/hifigan4')
hifi5_model = tf.keras.models.load_model('tfmodel/hifigan5')
def voc_infer(post_mel):
    hifi_in = tf.expand_dims(tf.convert_to_tensor(post_mel),axis=0)
    hifi0_out,exci_out = hifi0_model((hifi_in),training=False)
    print('hifi0 {} {}'.format(hifi0_out.shape,exci_out.shape))

    hifi1_win = 21
    hifi1_olp = 3
    hifi1_up = 10
    hifi_out_list = []
    time_len = tf.shape(hifi0_out).numpy()[1]
    for step in range(hifi1_olp,time_len,hifi1_win-2*hifi1_olp):
        start = step-hifi1_olp
        end = step+hifi1_win-hifi1_olp
        melstep = hifi0_out[:,start:end,:]
        excistep = exci_out[:,start*200:end*200]
        rpad = hifi1_win - tf.shape(melstep).numpy()[1]
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*200,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi1_model((melstep,excistep),training=True)
        #print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi1_olp:
            hifi_out_step = hifi_out_step[:,:-30,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,30:-30,:]
            else:
                if rpad > 0:
                    hifi_out_step = hifi_out_step[:,30:-(rpad*10),:]
                else:
                    hifi_out_step = hifi_out_step[:,30:,:]

        hifi_out_list.append(hifi_out_step)
        if end > time_len:
            break

    hifi1_out = tf.concat(hifi_out_list,axis=1)
    print('hifi1 {}'.format(hifi1_out.shape))

    hifi2_win = 150
    hifi2_olp = 5
    hifi2_up = 5
    hifi_out_list = []
    time_len = tf.shape(hifi1_out).numpy()[1]
    for step in range(hifi2_olp,time_len,hifi2_win-2*hifi2_olp):
        start = step-hifi2_olp
        end = step+hifi2_win-hifi2_olp
        melstep = hifi1_out[:,start:end,:]
        excistep = exci_out[:,(start)*20:(end)*20]
        rpad = hifi2_win - tf.shape(melstep).numpy()[1]
        #print('rpad {} {} {} {}'.format(step,rpad,melstep.shape,excistep.shape))
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*20,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi2_model((melstep,excistep),training=True)
        #print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi2_olp:
            hifi_out_step = hifi_out_step[:,:-25,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,25:-25,:]
            else:
                if rpad > 0:
                    hifi_out_step = hifi_out_step[:,25:-(rpad*5),:]
                else:
                    hifi_out_step = hifi_out_step[:,25:,:]
        hifi_out_list.append(hifi_out_step)
        if end > time_len:
            break

    hifi2_out = tf.concat(hifi_out_list,axis=1)
    print('hifi2 {}'.format(hifi2_out.shape))

    hifi3_win = 700
    hifi3_olp = 20
    hifi3_up = 2
    hifi_out_list = []
    time_len = tf.shape(hifi2_out).numpy()[1]
    for step in range(hifi3_olp,time_len,hifi3_win-2*hifi3_olp):
        start = step-hifi3_olp
        end = step+hifi3_win-hifi3_olp
        melstep = hifi2_out[:,start:end,:]
        excistep = exci_out[:,start*4:end*4]
        rpad = hifi3_win - tf.shape(melstep).numpy()[1]
        #print('rpad {} {} {}'.format(rpad,melstep.shape,excistep.shape))
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*4,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi3_model((melstep,excistep),training=True)
        #print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi3_olp:
            hifi_out_step = hifi_out_step[:,:-40,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,40:-40,:]
            else:
                if rpad > 0:
                    hifi_out_step = hifi_out_step[:,40:-(rpad*2),:]
                else:
                    hifi_out_step = hifi_out_step[:,40:,:]
        hifi_out_list.append(hifi_out_step)
        if end > time_len:
            break

    hifi3_out = tf.concat(hifi_out_list,axis=1)
    print('hifi3 {}'.format(hifi3_out.shape))

    hifi4_win = 1320 
    hifi4_olp = 20
    hifi4_up = 2
    hifi_out_list = []
    time_len = tf.shape(hifi3_out).numpy()[1]
    for step in range(hifi4_olp,time_len,hifi4_win-2*hifi4_olp):
        start = step-hifi4_olp
        end = step+hifi4_win-hifi4_olp
        melstep = hifi3_out[:,start:end,:]
        excistep = exci_out[:,(start)*2:(end)*2]
        rpad = hifi4_win - tf.shape(melstep).numpy()[1]
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
            excistep = tf.concat([excistep,tf.zeros([1,rpad*2,excistep.shape[2]])],axis=1)
        hifi_out_step = hifi4_model((melstep,excistep),training=True)
        #print('hifi_out {} '.format(hifi_out_step.shape))
        if step == hifi4_olp:
            hifi_out_step = hifi_out_step[:,:-40,:]
        else:
            if end <= time_len:
                hifi_out_step = hifi_out_step[:,40:-40,:]
            else:
                if rpad > 0:
                    hifi_out_step = hifi_out_step[:,40:-(rpad*2),:]
                else:
                    hifi_out_step = hifi_out_step[:,40:,:]
        hifi_out_list.append(hifi_out_step)
        if end > time_len:
            break

    hifi4_out = tf.concat(hifi_out_list,axis=1)
    print('hifi4 {}'.format(hifi4_out.shape))

    hifi5_win = 8008 
    hifi5_olp = 4
    hifi_out_list = []
    time_len = tf.shape(hifi4_out).numpy()[1]
    for step in range(4,time_len,8000):
        start = step-4
        end = step+hifi5_win-hifi5_olp
        melstep = hifi4_out[:,step-4:step+8004,:]
        rpad = 8008 - tf.shape(melstep).numpy()[1]
        if rpad > 0:
            melstep = tf.concat([melstep,tf.zeros([1,rpad,melstep.shape[2]])],axis=1)
        hifi_out_step = hifi5_model((melstep),training=True)
        #print('hifi_out {} '.format(hifi_out_step.shape))
        if step == 4:
            hifi_out_step = hifi_out_step[:,:-4,:]
        else:
            if end<=time_len:
                hifi_out_step = hifi_out_step[:,4:-4,:]
            else:
                if rpad > 0:
                    hifi_out_step = hifi_out_step[:,4:-(rpad),:]
                else:
                    hifi_out_step = hifi_out_step[:,4:,:]
        hifi_out_list.append(hifi_out_step)
        if end > time_len:
            break

    hifi5_out = tf.concat(hifi_out_list,axis=1)
    print('hifi5 {}'.format(hifi5_out.shape))
    pcm = (tf.reshape(hifi5_out,(-1))*32768.0).numpy().astype(dtype=np.int16)
    return pcm
    #import wave
    #with wave.open("tfdebug_{}.wav".format(fileid), 'wb') as wf:
    #    wf.setnchannels(1)
    #    wf.setsampwidth(2)
    #    wf.setframerate(16000)
    #    wf.writeframes(pcm.data)
